import torch
import torch.nn as nn
import torch.nn.functional as F


def get_cosine_loss(pred_vector, target_vector, reduction=True):
    cosine_vector_loss_dict = {}
    #  normalize pred_vector to unit vectors
    normalized_pred_vector = F.normalize(pred_vector)
    normalized_target_vector = F.normalize(target_vector)
    loss = 1 - normalized_pred_vector.mul(normalized_target_vector).sum(dim=1)
    if reduction:
        losses = loss.mean()
    else:
        losses = loss.sum()
    cosine_vector_loss_dict['consine_vector_loss'] = losses.item()
    return losses, cosine_vector_loss_dict



